{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff31084d",
   "metadata": {},
   "source": [
    "# Entity Memory\n",
    "This notebook shows how to work with a memory module that remembers things about specific entities. It extracts information on entities (using LLMs) and builds up its knowledge about that entity over time (also using LLMs).\n",
    "\n",
    "Let's first walk through using this functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1bea1181",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI\n",
    "from langchain.memory import ConversationEntityMemory\n",
    "llm = OpenAI(temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "34425079",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = ConversationEntityMemory(llm=llm)\n",
    "_input = {\"input\": \"Deven & Sam are working on a hackathon project\"}\n",
    "memory.load_memory_variables(_input)\n",
    "memory.save_context(\n",
    "    _input,\n",
    "    {\"ouput\": \" That sounds like a great project! What kind of project are they working on?\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b425642c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'Human: Deven & Sam are working on a hackathon project\\nAI:  That sounds like a great project! What kind of project are they working on?',\n",
       " 'entities': {'Sam': 'Sam is working on a hackathon project with Deven.'}}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({\"input\": 'who is Sam'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3bf89b46",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = ConversationEntityMemory(llm=llm, return_messages=True)\n",
    "_input = {\"input\": \"Deven & Sam are working on a hackathon project\"}\n",
    "memory.load_memory_variables(_input)\n",
    "memory.save_context(\n",
    "    _input,\n",
    "    {\"ouput\": \" That sounds like a great project! What kind of project are they working on?\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3e37d126",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [HumanMessage(content='Deven & Sam are working on a hackathon project', additional_kwargs={}),\n",
       "  AIMessage(content=' That sounds like a great project! What kind of project are they working on?', additional_kwargs={})],\n",
       " 'entities': {'Sam': 'Sam is working on a hackathon project with Deven.'}}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({\"input\": 'who is Sam'})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5ad043",
   "metadata": {},
   "source": [
    "## Using in a chain\n",
    "Let's now use it in a chain!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "13471fbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import ConversationChain\n",
    "from langchain.memory import ConversationEntityMemory\n",
    "from langchain.memory.prompt import ENTITY_MEMORY_CONVERSATION_TEMPLATE\n",
    "from pydantic import BaseModel\n",
    "from typing import List, Dict, Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "183346e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True,\n",
    "    prompt=ENTITY_MEMORY_CONVERSATION_TEMPLATE,\n",
    "    memory=ConversationEntityMemory(llm=llm)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7eb1460a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Deven': '', 'Sam': ''}\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Last line:\n",
      "Human: Deven & Sam are working on a hackathon project\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' That sounds like a great project! What kind of project are they working on?'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Deven & Sam are working on a hackathon project\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0269f513",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Deven': 'Deven is working on a hackathon project with Sam.',\n",
       " 'Sam': 'Sam is working on a hackathon project with Deven.'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.memory.store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "46324ca8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Deven': 'Deven is working on a hackathon project with Sam.', 'Sam': 'Sam is working on a hackathon project with Deven.', 'Langchain': ''}\n",
      "\n",
      "Current conversation:\n",
      "Human: Deven & Sam are working on a hackathon project\n",
      "AI:  That sounds like a great project! What kind of project are they working on?\n",
      "Last line:\n",
      "Human: They are trying to add more complex memory structures to Langchain\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' That sounds like an interesting project! What kind of memory structures are they trying to add?'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"They are trying to add more complex memory structures to Langchain\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ff2ebf6b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Deven': 'Deven is working on a hackathon project with Sam, attempting to add more complex memory structures to Langchain.', 'Sam': 'Sam is working on a hackathon project with Deven, trying to add more complex memory structures to Langchain.', 'Langchain': 'Langchain is a project that is trying to add more complex memory structures.', 'Key-Value Store': ''}\n",
      "\n",
      "Current conversation:\n",
      "Human: Deven & Sam are working on a hackathon project\n",
      "AI:  That sounds like a great project! What kind of project are they working on?\n",
      "Human: They are trying to add more complex memory structures to Langchain\n",
      "AI:  That sounds like an interesting project! What kind of memory structures are they trying to add?\n",
      "Last line:\n",
      "Human: They are adding in a key-value store for entities mentioned so far in the conversation.\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' That sounds like a great idea! How will the key-value store work?'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"They are adding in a key-value store for entities mentioned so far in the conversation.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "56cfd4ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Deven': 'Deven is working on a hackathon project with Sam, attempting to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation.', 'Sam': 'Sam is working on a hackathon project with Deven, trying to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation.'}\n",
      "\n",
      "Current conversation:\n",
      "Human: Deven & Sam are working on a hackathon project\n",
      "AI:  That sounds like a great project! What kind of project are they working on?\n",
      "Human: They are trying to add more complex memory structures to Langchain\n",
      "AI:  That sounds like an interesting project! What kind of memory structures are they trying to add?\n",
      "Human: They are adding in a key-value store for entities mentioned so far in the conversation.\n",
      "AI:  That sounds like a great idea! How will the key-value store work?\n",
      "Last line:\n",
      "Human: What do you know about Deven & Sam?\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Deven and Sam are working on a hackathon project together, attempting to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation.'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"What do you know about Deven & Sam?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e6df549",
   "metadata": {},
   "source": [
    "## Inspecting the memory store\n",
    "We can also inspect the memory store directly. In the following examaples, we look at it directly, and then go through some examples of adding information and watch how it changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "038b4d3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Deven': 'Deven is working on a hackathon project with Sam, attempting to add '\n",
      "          'more complex memory structures to Langchain, including a key-value '\n",
      "          'store for entities mentioned so far in the conversation.',\n",
      " 'Key-Value Store': 'A key-value store that stores entities mentioned in the '\n",
      "                    'conversation.',\n",
      " 'Langchain': 'Langchain is a project that is trying to add more complex '\n",
      "              'memory structures, including a key-value store for entities '\n",
      "              'mentioned so far in the conversation.',\n",
      " 'Sam': 'Sam is working on a hackathon project with Deven, attempting to add '\n",
      "        'more complex memory structures to Langchain, including a key-value '\n",
      "        'store for entities mentioned so far in the conversation.'}\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "pprint(conversation.memory.store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2df4800e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Daimon': '', 'Sam': 'Sam is working on a hackathon project with Deven to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation.'}\n",
      "\n",
      "Current conversation:\n",
      "Human: They are trying to add more complex memory structures to Langchain\n",
      "AI:  That sounds like an interesting project! What kind of memory structures are they trying to add?\n",
      "Human: They are adding in a key-value store for entities mentioned so far in the conversation.\n",
      "AI:  That sounds like a great idea! How will the key-value store work?\n",
      "Human: What do you know about Deven & Sam?\n",
      "AI:  Deven and Sam are working on a hackathon project to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation. They seem to be very motivated and passionate about their project, and are working hard to make it a success.\n",
      "Last line:\n",
      "Human: Sam is the founder of a company called Daimon.\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"\\nThat's impressive! It sounds like Sam is a very successful entrepreneur. What kind of company is Daimon?\""
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Sam is the founder of a company called Daimon.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ebe9e36f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'Daimon': 'Daimon is a company founded by Sam.',\n",
      " 'Deven': 'Deven is working on a hackathon project with Sam to add more '\n",
      "          'complex memory structures to Langchain, including a key-value store '\n",
      "          'for entities mentioned so far in the conversation.',\n",
      " 'Key-Value Store': 'Key-Value Store: A data structure that stores values '\n",
      "                    'associated with a unique key, allowing for efficient '\n",
      "                    'retrieval of values. Deven and Sam are adding a key-value '\n",
      "                    'store for entities mentioned so far in the conversation.',\n",
      " 'Langchain': 'Langchain is a project that seeks to add more complex memory '\n",
      "              'structures, including a key-value store for entities mentioned '\n",
      "              'so far in the conversation.',\n",
      " 'Sam': 'Sam is working on a hackathon project with Deven to add more complex '\n",
      "        'memory structures to Langchain, including a key-value store for '\n",
      "        'entities mentioned so far in the conversation. He is also the founder '\n",
      "        'of a company called Daimon.'}\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "pprint(conversation.memory.store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dd547144",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mYou are an assistant to a human, powered by a large language model trained by OpenAI.\n",
      "\n",
      "You are designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, you are able to generate human-like text based on the input you receive, allowing you to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.\n",
      "\n",
      "You are constantly learning and improving, and your capabilities are constantly evolving. You are able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. You have access to some personalized information provided by the human in the Context section below. Additionally, you are able to generate your own text based on the input you receive, allowing you to engage in discussions and provide explanations and descriptions on a wide range of topics.\n",
      "\n",
      "Overall, you are a powerful tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether the human needs help with a specific question or just wants to have a conversation about a particular topic, you are here to assist.\n",
      "\n",
      "Context:\n",
      "{'Sam': 'Sam is working on a hackathon project with Deven to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation. He is also the founder of a company called Daimon.', 'Daimon': 'Daimon is a company founded by Sam.'}\n",
      "\n",
      "Current conversation:\n",
      "Human: They are adding in a key-value store for entities mentioned so far in the conversation.\n",
      "AI:  That sounds like a great idea! How will the key-value store work?\n",
      "Human: What do you know about Deven & Sam?\n",
      "AI:  Deven and Sam are working on a hackathon project to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation. They seem to be very motivated and passionate about their project, and are working hard to make it a success.\n",
      "Human: Sam is the founder of a company called Daimon.\n",
      "AI: \n",
      "That's impressive! It sounds like Sam is a very successful entrepreneur. What kind of company is Daimon?\n",
      "Last line:\n",
      "Human: What do you know about Sam?\n",
      "You:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Sam is the founder of a company called Daimon. He is also working on a hackathon project with Deven to add more complex memory structures to Langchain, including a key-value store for entities mentioned so far in the conversation. He seems to be very motivated and passionate about his project, and is working hard to make it a success.'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"What do you know about Sam?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e00463b5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
